﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;

namespace SkypeMonitorService
{
	internal class TerminalServiceManager
	{
		public IntPtr WTS_CURRENT_SERVER_HANDLE = IntPtr.Zero;
		public const int WTS_CURRENT_SESSION = -1;

		[DllImport("wtsapi32.dll", SetLastError = true)]
		static extern IntPtr WTSOpenServer([MarshalAs(UnmanagedType.LPStr)] String pServerName);

		[DllImport("wtsapi32.dll")]
		static extern void WTSCloseServer(IntPtr hServer);

		[DllImport("wtsapi32.dll", SetLastError = true)]
		static extern Int32 WTSEnumerateSessions(
			IntPtr hServer,
			[MarshalAs(UnmanagedType.U4)] Int32 Reserved,
			[MarshalAs(UnmanagedType.U4)] Int32 Version,
			ref IntPtr ppSessionInfo,
			[MarshalAs(UnmanagedType.U4)] ref Int32 pCount);

		[DllImport("wtsapi32.dll")]
		static extern void WTSFreeMemory(IntPtr pMemory);

		[DllImport("wtsapi32.dll", SetLastError = true)]
		static extern bool WTSSendMessage(IntPtr hServer, [MarshalAs(UnmanagedType.I4)] int SessionId,
					 String pTitle,
					 [MarshalAs(UnmanagedType.U4)] int TitleLength,
					 String pMessage,
					 [MarshalAs(UnmanagedType.U4)] int MessageLength,
					 [MarshalAs(UnmanagedType.U4)] int Style,
					 [MarshalAs(UnmanagedType.U4)] int Timeout,
					 [MarshalAs(UnmanagedType.U4)] out int pResponse,
					 bool bWait);

		[DllImport("Kernel32.dll", SetLastError = true)]
		static extern int WTSGetActiveConsoleSessionId();

		[DllImport("wtsapi32.dll", SetLastError = true)]
		static extern bool WTSLogoffSession(IntPtr hServer, int SessionId, bool bWait);

		[DllImport("wtsapi32.dll", SetLastError = true)]
		static extern bool WTSQueryUserToken(int SessionId, out  IntPtr phToken);

		[DllImport("wtsapi32.dll", SetLastError = true)]
		static extern bool WTSEnumerateServers(string pDomainName, uint Reserved, uint Version, ref IntPtr ppServerInfo, ref uint pCount);

		[StructLayout(LayoutKind.Sequential)]
		private struct WTS_SESSION_INFO
		{
			public Int32 SessionID;

			[MarshalAs(UnmanagedType.LPStr)]
			public String pWinStationName;

			public WTS_CONNECTSTATE_CLASS State;
		}

		public enum WTS_CONNECTSTATE_CLASS
		{
			WTSActive,
			WTSConnected,
			WTSConnectQuery,
			WTSShadow,
			WTSDisconnected,
			WTSIdle,
			WTSListen,
			WTSReset,
			WTSDown,
			WTSInit
		}

		[StructLayout(LayoutKind.Sequential)]
		internal struct WTS_SERVER_INFO
		{
			public string pServerName;
		}

		public static IntPtr OpenServer(String Name)
		{
			IntPtr server = WTSOpenServer(Name);
			return server;
		}

		public static void CloseServer(IntPtr ServerHandle)
		{
			WTSCloseServer(ServerHandle);
		}

		public static List<String> ListSessions(String ServerName)
		{
			IntPtr server = IntPtr.Zero;
			List<String> ret = new List<string>();
			server = OpenServer(ServerName);

			try
			{
				IntPtr ppSessionInfo = IntPtr.Zero;

				Int32 count = 0;
				Int32 retval = WTSEnumerateSessions(server, 0, 1, ref ppSessionInfo, ref count);
				Int32 dataSize = Marshal.SizeOf(typeof(WTS_SESSION_INFO));

				Int64 current = (int)ppSessionInfo;

				if (retval != 0)
				{
					for (int i = 0; i < count; i++)
					{
						WTS_SESSION_INFO si = (WTS_SESSION_INFO)Marshal.PtrToStructure((System.IntPtr)current, typeof(WTS_SESSION_INFO));
						current += dataSize;

						ret.Add(si.SessionID + " " + si.State + " " + si.pWinStationName);
					}

					WTSFreeMemory(ppSessionInfo);
				}
			}
			finally
			{
				CloseServer(server);
			}

			return ret;
		}

		/// <summary>
		/// Send terminal service message
		/// </summary>
		/// <exception cref="System.ComponentModel.Win32Exception"/>
		/// <param name="machineName"></param>
		/// <param name="msgs"></param>
		/// <param name="title"></param>
		public static void SendMessage(String machineName, String msgs, String title)
		{
			IntPtr server = IntPtr.Zero;
			server = OpenServer(machineName);
			int resp = 0;
			try
			{
				if (WTSSendMessage(server, WTSGetActiveConsoleSessionId(), title, title.Length, msgs, msgs.Length, 0, 0, out resp, false))
				{
				}
				else
				{
					int err = Marshal.GetLastWin32Error();
					throw new System.ComponentModel.Win32Exception(err);
				}
			}
			finally
			{
				if (server != IntPtr.Zero)
					CloseServer(server);
			}
		}


		/// <summary>
		/// Enumerate avaliable servers
		/// </summary>
		/// <param name="domainName">Current domain name (can be null if you want to enumerate for current)</param>
		/// <returns></returns>
		public static List<string> EnumerateServers(string domainName)
		{
			IntPtr pInfo = IntPtr.Zero;
			uint count = 0;
			Int32 dataSize = Marshal.SizeOf(typeof(WTS_SERVER_INFO));
			List<string> result = new List<string>();

			if (WTSEnumerateServers(domainName, 0, 1, ref pInfo, ref count))
			{
				Int64 current = (int)pInfo;
				for (int i = 0; i < count; i++)
				{
					WTS_SERVER_INFO WTScsi = (WTS_SERVER_INFO)Marshal.PtrToStructure((IntPtr)current, typeof(WTS_SERVER_INFO));
					current += dataSize;
					result.Add(WTScsi.pServerName);
				}
				WTSFreeMemory(pInfo);
				return result;
			}
			else
			{
				throw new System.ComponentModel.Win32Exception(Marshal.GetLastWin32Error());
			}
			return null;
		}
	}
}
